import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from sklearn.metrics import f1_score,accuracy_score,classification_report,confusion_matrix
import time, os
from utils_common import get_time_dif, build_class_map, load_MLB
from tensorboardX import SummaryWriter
from utils_bert import freeze_to_layer_by_name
from transformers import AdamW, get_linear_schedule_with_warmup


# 权重初始化，默认xavier
def init_network(model, method='xavier', exclude='embedding', seed=123, args=None):
    for name, w in model.named_parameters():
        if exclude not in name:
            if args and "bert" in args.model and len(w.size()) < 2:
                # why?
                continue
            if 'weight' in name:
                if method == 'xavier':
                    nn.init.xavier_normal_(w)
                elif method == 'kaiming':
                    nn.init.kaiming_normal_(w)
                else:
                    nn.init.normal_(w)
            elif 'bias' in name:
                nn.init.constant_(w, 0)
            else:
                pass


## ----------------------single_label----------------------------- ##

def train(config, model, train_iter, dev_iter, test_iter, args=None):
    """单分类
    """
    start_time = time.time()

    if "bert" in args.model:
        param_optimizer = list(model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]
        
        max_grad_norm = 1.0
        num_training_steps = len(train_iter) * config.num_epochs
        num_warmup_steps = int(0.1*(num_training_steps))
        optimizer = AdamW(optimizer_grouped_parameters, lr=config.learning_rate, correct_bias=False)
        scheduler = get_linear_schedule_with_warmup(optimizer, 
                        num_warmup_steps=num_warmup_steps, 
                        num_training_steps=num_training_steps)
        freeze_to_layer_by_name(model, args.freeze_to_layer)
    else:
        optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
        scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)

    total_batch = 0  # 记录进行到多少batch
    dev_best_loss = float('inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升
    writer = SummaryWriter(logdir=config.log_path + '/' + time.strftime('%m-%d_%H.%M', time.localtime()))
    model.train()
    for epoch in range(config.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
        for i, something in enumerate(train_iter):
            if "bert" in args.model:
                labels = something["labels"]
                trains = something
            else:
                trains, labels = something
            outputs = model(trains)
            model.zero_grad()
            loss = F.cross_entropy(outputs, labels)
            loss.backward()
            optimizer.step()
            scheduler.step()
            if total_batch % 100 == 0:
                true = labels.data.cpu()
                predic = torch.max(outputs.data, 1)[1].cpu()
                train_acc = accuracy_score(true, predic)
                dev_acc, dev_loss = evaluate(config, model, dev_iter)
                if dev_loss < dev_best_loss:
                    dev_best_loss = dev_loss
                    torch.save(model.state_dict(), config.save_path)
                    improve = '*'
                    last_improve = total_batch
                else:
                    improve = ''
                time_dif = get_time_dif(start_time)
                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Train Acc: {2:>6.2%},  Val Loss: {3:>5.2},  Val Acc: {4:>6.2%},  Time: {5} {6}'
                print(msg.format(total_batch, loss.item(), train_acc, dev_loss, dev_acc, time_dif, improve))
                writer.add_scalar("loss/train", loss.item(), total_batch)
                writer.add_scalar("loss/dev", dev_loss, total_batch)
                writer.add_scalar("acc/train", train_acc, total_batch)
                writer.add_scalar("acc/dev", dev_acc, total_batch)
                model.train()
            total_batch += 1
            if total_batch - last_improve > config.require_improvement:
                print("No optimization for a long time, auto-stopping...")
                flag = True
                break
        if flag:
            break
    writer.close()
    test(config, model, test_iter)

def test(config, model, test_iter):
    model.load_state_dict(torch.load(config.save_path))
    model.eval()
    start_time = time.time()
    test_acc, test_loss, test_report, test_confusion = evaluate(config, model, test_iter, test=True)
    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}'
    print(msg.format(test_loss, test_acc))
    print("Precision, Recall and F1-Score...")
    print(test_report)
    print("Confusion Matrix...")
    print(test_confusion)
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)

def evaluate(config, model, data_iter, test=False):
    model.eval()
    loss_total = 0
    predicts_all = np.array([], dtype=int)
    labels_all = np.array([], dtype=int)
    with torch.no_grad():
        for texts, labels in data_iter:
            outputs = model(texts)
            loss = F.cross_entropy(outputs, labels)
            loss_total += loss
            labels = labels.data.cpu().numpy()
            predic = torch.max(outputs.data, 1)[1].cpu().numpy()
            labels_all = np.append(labels_all, labels)
            predicts_all = np.append(predicts_all, predic)

    acc = accuracy_score(labels_all, predicts_all)
    if test:
        labels_all=labels_all.tolist()
        predicts_all=predicts_all.tolist()
        try:
            # 处理异常: labels_all的长度和config.class_list不一致会报错
            if len(set(labels_all)) != len(config.class_list):
                # 找到缺失的idx, 补到pred里面, 算预测对
                no_find_idx = [i for i in range(len(config.class_list)) if i not in labels_all]
                labels_all.extend(no_find_idx)
                predicts_all.extend(no_find_idx)
                print(f"NO FOUND label: {' '.join([config.class_list[idx] for idx in no_find_idx])}")

            report = classification_report(labels_all, predicts_all, target_names=config.class_list, digits=4)
            confusion = confusion_matrix(labels_all, predicts_all)
        except ValueError:
            report = confusion = "ValueError in metrics report"
        return acc, loss_total / len(data_iter), report, confusion
    return acc, loss_total / len(data_iter)

def infer(config, model, data_iter, args):
    """单分类
    """
    model.eval()
    predicts_all = []
    with torch.no_grad():
        for texts, _labels in data_iter:
            outputs = model(texts)
            predic = torch.max(outputs.data, 1)[1].cpu().numpy()
            predicts_all += predic.tolist()

    if args.infer_combine:
        with open(config.infer_path, 'r') as f1:
            texts_all = f1.readlines()
        texts_all = [line.strip() for line in texts_all if line.strip()]
        assert len(texts_all) == len(predicts_all)
        
        _class_to_idx, idx_to_class = build_class_map(config)
        predicts_all_names = [idx_to_class[idx] for idx in predicts_all]
        with open(f"{config.infer_path[:-4]}_{config.model_name}_with_label.txt", 'w') as f1:
            for text, name in zip(texts_all, predicts_all_names):
                f1.write(f"{text}\t{name}\n")
        print(f"Write to file: {config.infer_path[:-4]}_{config.model_name}_with_label.txt")
    else:
        with open(f"{config.infer_path[:-4]}_{config.model_name}_only_label.txt", 'w') as f1:
            for name in predicts_all_names:
                f1.write(f"{name}\n")
        print(f"Write to file: {config.infer_path[:-4]}_{config.model_name}_only_label.txt")


## ----------------------multi_labels----------------------------- ##


def train_multi_labels(config, model, train_iter, dev_iter, test_iter, args):
    start_time = time.time()

    if "bert" in args.model:
        param_optimizer = list(model.named_parameters())
        no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
        optimizer_grouped_parameters = [
            {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
            {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]
        
        max_grad_norm = 1.0
        num_training_steps = len(train_iter) * config.num_epochs
        num_warmup_steps = int(0.1*(num_training_steps))
        optimizer = AdamW(optimizer_grouped_parameters, lr=config.learning_rate, correct_bias=False)
        scheduler = get_linear_schedule_with_warmup(optimizer, 
                        num_warmup_steps=num_warmup_steps, 
                        num_training_steps=num_training_steps)
        freeze_to_layer_by_name(model, args.freeze_to_layer)
    else:
        optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
        scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)

    # criterion = nn.BCEWithLogitsLoss(pos_weight=config.class_weights)
    criterion = nn.BCEWithLogitsLoss()
    total_batch = 0  # 记录进行到多少batch
    dev_best_loss = float('inf')
    last_improve = 0  # 记录上次验证集loss下降的batch数
    flag = False  # 记录是否很久没有效果提升
    writer = SummaryWriter(logdir=config.log_path + '/' + time.strftime('%m-%d_%H.%M', time.localtime()))
    model.train()
    for epoch in range(config.num_epochs):
        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))
        for i, something in enumerate(train_iter):
            if "bert" in args.model:
                labels = something["labels"]
                trains = something
            else:
                trains, labels = something
            outputs = model(trains)
            model.zero_grad()
            loss = criterion(outputs, labels.float())
            loss.backward()
            optimizer.step()
            scheduler.step() # 学习率衰减
            if total_batch % args.print_per_batch == 0:
                dev_acc, dev_f1_score_macro, dev_f1_score_micro, dev_loss = evaluate_multi_labels(config, model, dev_iter, args)
                if dev_loss < dev_best_loss:
                    dev_best_loss = dev_loss
                    torch.save(model.state_dict(), config.save_path)
                    improve = '*'
                    last_improve = total_batch
                else:
                    improve = ''
                time_dif = get_time_dif(start_time)
                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Val Macro F1: {2:>6.2%},  Val Micro F1: {3:>6.2%}, Val Loss: {4:>5.2},  Val Acc: {5:>6.2%},  Time: {6} {7}'
                print(msg.format(total_batch, loss.item(), dev_f1_score_macro, dev_f1_score_micro, dev_loss, dev_acc, time_dif, improve))
                writer.add_scalar("loss/train", loss.item(), total_batch)
                writer.add_scalar("loss/dev", dev_loss, total_batch)
                writer.add_scalar("acc/dev", dev_acc, total_batch)
                writer.add_scalar("Macro F1/dev", dev_f1_score_macro, total_batch)
                writer.add_scalar("Micro F1/dev", dev_f1_score_micro, total_batch)
                model.train()
            total_batch += 1
            if total_batch - last_improve > config.require_improvement:
                print("No optimization for a long time, auto-stopping...")
                flag = True
                break
        if flag:
            break
    writer.close()
    test_multi_labels(config, model, test_iter, args)

def test_multi_labels(config, model, test_iter, args):
    # test
    model.load_state_dict(torch.load(config.save_path))
    model.eval()
    start_time = time.time()
    test_acc, f1_score_macro, f1_score_micro, test_loss = evaluate_multi_labels(config, model, test_iter, args)
    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}, Test Micro F1: {2:>6.2%}, Test Micro F1: {3:>6.2%}'
    print(msg.format(test_loss, test_acc, f1_score_macro, f1_score_micro))
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)

def evaluate_multi_labels(config, model, data_iter, args):
    """改写为多标签.
    """
    model.eval()
    loss_total = 0
    labels_all = []
    predicts_all = []
    criterion = nn.BCEWithLogitsLoss()
    with torch.no_grad():
        for i, something in enumerate(data_iter):
            if "bert" in args.model:
                labels = something["labels"]
                texts = something
            else:
                texts, labels = something
            outputs = model(texts)
            loss = criterion(outputs, labels.float())
            loss_total += loss
            labels = labels.data.cpu().numpy()
            outputs = outputs.data.cpu().numpy()
            predicts = np.where(outputs > args.threshold, 1, 0)
            labels_all += labels.tolist()
            predicts_all += predicts.tolist()

    labels_all = np.array(labels_all,dtype=int) 
    predicts_all = np.array(predicts_all,dtype=int)

    f1_score_macro = f1_score(labels_all, predicts_all, average='macro')
    f1_score_micro = f1_score(labels_all, predicts_all, average='micro')
    accuracy = accuracy_score(labels_all, predicts_all)
    test_loss = loss_total / len(data_iter)

    return accuracy, f1_score_macro, f1_score_micro, test_loss

def infer_multi_labels(config, model, data_iter, args):
    """多标签分类
    """
    model.eval()
    mlb = load_MLB(config)
    
    predicts_all = []
    with torch.no_grad():
        for i, something in enumerate(train_iter):
            if "bert" in args.model:
                labels = something["labels"]
                texts = something
            else:
                texts, labels = something
            outputs = model(texts)
            outputs = outputs.data.cpu().numpy()
            predicts = np.where(outputs > args.threshold, 1, 0)
            predicts_all += predicts.tolist()
    predicts_all = np.array(predicts_all, dtype=int)
    predicts_all_names = mlb.inverse_transform(predicts_all)

    if args.infer_combine:
        with open(config.infer_path, 'r') as f1:
            texts_all = f1.readlines()
        texts_all = [line.strip() for line in texts_all if line.strip()]
        assert len(texts_all) == len(predicts_all_names)
        
        with open(f"{config.infer_path[:-4]}_{config.model_name}_with_labels.txt", 'w') as f1:
            for text, names in zip(texts_all, predicts_all_names):
                f1.write(f"{text}\t" + "\t".join(names) + "\n")
        print(f"Write to file: {config.infer_path[:-4]}_{config.model_name}_with_labels.txt")
    else:
        with open(f"{config.infer_path[:-4]}_{config.model_name}_only_labels.txt", 'w') as f1:
            for names in predicts_all_names:
                f1.write("\t".joint(names) +"\n")
        print(f"Write to file: {config.infer_path[:-4]}_{config.model_name}_only_labels.txt")

